/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2020-2022. All rights reserved.
 * Description: 集合通信算子图优化
 * Author: jiangchen
 * Create: 2020-06-08
 */

#include "hvd_graph_optimizer.h"
#include "hvd_all_reduce_fusion.h"
#include "hcom_ops_kernel_info_store.h"
#include "graph/debug/ge_attr_define.h"
#include "comm.h"
using namespace std;

namespace hccl {
HvdGraphOptimizer::HvdGraphOptimizer()
{
}

HvdGraphOptimizer::~HvdGraphOptimizer()
{
}

ge::Status HvdGraphOptimizer::Initialize(const std::map<std::string, std::string> &options,
                                         ge::OptimizeUtility *const optimizeUtility)
{
    HCCL_INFO("init hccl graph optimizer.");
    return ge::SUCCESS;
}

ge::Status HvdGraphOptimizer::Finalize()
{
    HCCL_INFO("finalize hccl graph optimizer.");
    return ge::SUCCESS;
}

HcclResult HvdGraphOptimizer::CheckSupportedOP(const std::string &sCollectiveType) const
{
    std::vector<std::string>::const_iterator it = std::find(HVD_SUPPORTED_OP_TYPE.begin(),
        HVD_SUPPORTED_OP_TYPE.end(), sCollectiveType);
    return (it != HVD_SUPPORTED_OP_TYPE.end()) ? HCCL_SUCCESS : HCCL_E_PARA;
}

ge::Status HvdGraphOptimizer::OptimizeOriginalGraph(ge::ComputeGraph& graph)
{
    HvdAllReduceFusion fusionHvdAllReduceOp;
    HcclResult ret = fusionHvdAllReduceOp.Run(graph);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("[Optimize][OriginalGraph]graph[%s]: fuse HvdAllReduce op failed. ret[%d]",
            graph.GetName().c_str(), ret), ge::INTERNAL_ERROR);

    // insert hvdWait op to graph
    u32 eventCount = 0;
    std::string event = "event";
    std::vector<ge::NodePtr> addedWaitNode;
    HCCL_INFO("start traverse graph ans insert hvdWait.");
    u32 opCount = 0;
    std::vector<ge::NodePtr> broadcastOps;
    for (auto nodePtr : graph.GetDirectNode()) {
        if (!nodePtr) {
            HCCL_WARNING("null node exists.");
            continue;
        }
        auto opDescPtr = nodePtr->GetOpDesc();
        if (!opDescPtr) {
            HCCL_WARNING("desc of node is null.");
            continue;
        }
        if (opDescPtr->GetType() == HVD_KERNEL_OP_TYPE_BROADCAST || opDescPtr->GetType() == HVD_KERNEL_OP_TYPE_ALLREDUCE
            || opDescPtr->GetType() == HVD_KERNEL_OP_TYPE_ALLGATHER) {
            if (opDescPtr->GetType() == HVD_KERNEL_OP_TYPE_BROADCAST) {
                broadcastOps.push_back(nodePtr);
            }
            // the number of wait op and event is the same
            std::string tempOpName = "hvdWait" + std::to_string(eventCount);
            ge::OpDescPtr addedDescPtr = nullptr;
            EXECEPTION_CATCH((addedDescPtr = std::make_shared<ge::OpDesc>(tempOpName, HVD_KERNEL_OP_TYPE_WAIT)),
                return ge::FAILED);

            // set EVENT_IDENTIFY for both current node and wait node
            std::string tempevent = event + std::to_string(eventCount);
            bool setStrResult = (!ge::AttrUtils::SetStr(addedDescPtr, ATTR_INTER_EVENT_IDENTIFY, tempevent)) ||
                (!ge::AttrUtils::SetStr(opDescPtr, ATTR_INTER_EVENT_IDENTIFY, tempevent));
            CHK_PRT_RET(setStrResult, HCCL_ERROR("[Optimize][Graph]Set eventIdentify for addedNode/current"
                        "Node failed."), ge::FAILED);
            HCCL_DEBUG("set EVENT_IDENTIFY success, event: [%s]", tempevent.c_str());

            std::string waitStream = "waitStream";
            std::string sendStream = "sendStream";

            ret = AssignNewStream(opCount, nodePtr, waitStream, sendStream);
            CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Optimize][Graph]AssignNewStream failed"), ge::FAILED);

            CHK_PRT_RET(!ge::AttrUtils::SetStr(addedDescPtr, ATTR_NAME_STREAM_LABEL, waitStream),
                HCCL_ERROR("[Optimize][Graph]Set waitStream attr for added node failed"), ge::FAILED);

            CHK_PRT_RET(!ge::AttrUtils::SetStr(opDescPtr, ATTR_NAME_STREAM_LABEL, sendStream),
                HCCL_ERROR("[Optimize][Graph]Set sendStream attr for current node failed"), ge::FAILED);

            HCCL_DEBUG("set waitStream lable: [%s], sendStream lable: [%s]", waitStream.c_str(), sendStream.c_str());
            auto addedNodePtr = graph.AddNode(addedDescPtr);
            CHK_PRT_RET(!addedNodePtr, HCCL_ERROR("[Optimize][Graph]create hvdWait node failed"), ge::FAILED);

            addedWaitNode.push_back(addedNodePtr);
            ret = AddControlAnchor(nodePtr, addedNodePtr);
            CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Optimize][Graph]AddControlAnchor Failed."), ge::FAILED);
            eventCount++;
        }
    }
    // add control anchor between differnt hvdwait node
    // allgather和allreduce之间没控制边依赖，allreduce对应的wait和allgather对应的wait之间有控制边依赖
    ret = AddDependence(addedWaitNode);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Optimize][Graph]AddDependence wait failed"), ge::FAILED);

    ret = AddDependence(broadcastOps);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[Optimize][Graph]AddDependence broadcast failed"), ge::FAILED);

    HCCL_INFO("End OptimizeOriginalGraph");
    return ge::SUCCESS;
}

ge::Status HvdGraphOptimizer::OptimizeGraphPrepare(ge::ComputeGraph& graph)
{
    HCCL_INFO("start hccl graph optimizer prepare.");
    for (auto nodePtr : graph.GetAllNodes()) {
        if (!nodePtr) {
            HCCL_WARNING("OptimizeGraphPrepare: null node exists.");
            continue;
        }
        auto opDescPtr = nodePtr->GetOpDesc();
        if (!opDescPtr) {
            HCCL_WARNING("OptimizeGraphPrepare: desc of node[%s] is null.", nodePtr->GetName().c_str());
            continue;
        }
        // 集合通信的算子，设定format_Agnostic属性, 均设置成格式不敏感
        if (CheckSupportedOP(opDescPtr->GetType()) == HCCL_SUCCESS) {
            string name = "_format_agnostic";
            // 连续内存属性从图编译移至图优化
            bool bRet = ge::AttrUtils::SetBool(opDescPtr, ge::ATTR_NAME_CONTINUOUS_INPUT, true);
            CHK_PRT_RET(!bRet, HCCL_ERROR("[Optimize][GraphPrepare]errNo[0x%016llx] set continuous input attr[%d]"
                "to OpDesc failed.", HCOM_ERROR_CODE(HCCL_E_PARA), true), HCCL_E_PARA);
            bRet = ge::AttrUtils::SetBool(opDescPtr, ge::ATTR_NAME_CONTINUOUS_OUTPUT, true);
            CHK_PRT_RET(!bRet, HCCL_ERROR("[Optimize][GraphPrepare]errNo[0x%016llx] set continuous output attr[%d]"
                "to OpDesc failed.", HCOM_ERROR_CODE(HCCL_E_PARA), true), HCCL_E_PARA);

            bRet = ge::AttrUtils::SetInt(opDescPtr, name, HCCL_HVD_FORMAT_PAIRED_INPUT_OUTPUT);
            HCCL_DEBUG("graph[%s]: node [%s] op type [%s] format agnostic value is set",
                graph.GetName().c_str(), nodePtr->GetName().c_str(), opDescPtr->GetType().c_str());
            CHK_PRT_RET(!bRet, HCCL_ERROR("[Optimize][GraphPrepare]graph[%s]: node [%s] SetBool format_Agnostic"
                "failed, op type[%s]", graph.GetName().c_str(), nodePtr->GetName().c_str(),
                opDescPtr->GetType().c_str()), ge::INTERNAL_ERROR);
            if (opDescPtr->GetType() != HVD_KERNEL_OP_TYPE_BROADCAST) {
                name = "_input_mutable";
                bRet = ge::AttrUtils::SetBool(opDescPtr, name, true);
                HCCL_DEBUG("graph[%s]: node [%s] op type [%s] input mutable attr is set",
                    graph.GetName().c_str(), nodePtr->GetName().c_str(), opDescPtr->GetType().c_str());
                CHK_PRT_RET(!bRet, HCCL_ERROR("[Optimize][GraphPrepare]graph[%s]: node [%s]"
                    "SetBool _input_mutable failed, op type[%s]", graph.GetName().c_str(), nodePtr->GetName().c_str(),
                    opDescPtr->GetType().c_str()), ge::INTERNAL_ERROR);
            }
        }
    }
    HCCL_INFO("end hccl graph optimizer prepare.");
    return ge::SUCCESS;
}

HcclResult HvdGraphOptimizer::AssignNewStream(u32 &opCount, ge::NodePtr &nodePtr,
    std::string &waitStream, std::string &sendStream)
{
    auto opDescPtr = nodePtr->GetOpDesc();
    if (opDescPtr->GetType() == HVD_KERNEL_OP_TYPE_BROADCAST) {
        opCount++;
        std::string suffix = std::to_string(opCount / SPLIT_STREAM_THREAD);
        HCCL_DEBUG("Current satisfied node type is [%s], opCount [%u], suffix[%s] ",
            opDescPtr->GetType().c_str(), opCount, suffix.c_str());
        for (auto outDataAnchor : nodePtr->GetAllOutDataAnchors()) {
            if (!outDataAnchor) {
                continue;
            }
            std::string assign = "assign" + suffix;
            for (auto peerInDataAnchor : outDataAnchor->GetPeerInDataAnchors()) {
                if (!peerInDataAnchor) {
                    continue;
                }
                auto nextNodeDescPtr = peerInDataAnchor->GetOwnerNode()->GetOpDesc();
                if (nextNodeDescPtr->GetType() != "Assign") {
                    HCCL_ERROR("[Assign][NewStream]Not Assign after hvdBroadcast");
                    continue;
                }
                if (!ge::AttrUtils::SetStr(nextNodeDescPtr, ATTR_NAME_STREAM_LABEL, assign)) {
                    HCCL_ERROR("[Assign][NewStream]Set assign STREAM_LABEL failed");
                    return HCCL_E_INTERNAL;
                }
            }
            HCCL_DEBUG("Set assign STREAM_LABEL success [%s]", assign.c_str());
        }
        waitStream += suffix;
        sendStream += suffix;
    }
    return HCCL_SUCCESS;
}

// add control anchor between addedNode and downstream of current node
HcclResult HvdGraphOptimizer::AddControlAnchor(ge::NodePtr &nodePtr, ge::NodePtr &addedNodePtr)
{
    std::unordered_set<uintptr_t> anchorSet;
    anchorSet.clear();
    // 当前未考虑callback算子后面有纯控制边算子（非callback）
    // 遍历当前node的所有outdata anchor
    for (auto outDataAnchor : nodePtr->GetAllOutDataAnchors()) {
        if (!outDataAnchor) {
            continue;
        }
        // 对每个outdata anchor获取下游incontrol anchor，与wait算子的outcontrol连接
        for (auto peerInControlAnchor : outDataAnchor->GetPeerInControlAnchors()) {
            if (!peerInControlAnchor) {
                continue;
            }
            if (anchorSet.count((uintptr_t)peerInControlAnchor.get()) == 0) {
                ge::GraphUtils::AddEdge(addedNodePtr->GetOutControlAnchor(), peerInControlAnchor);
                anchorSet.insert((uintptr_t)peerInControlAnchor.get());
            }
        }
        HCCL_DEBUG("outDataAnchor linkto PeerInControlAnchor num is [%d]", anchorSet.size());
        // 对每个outdata anchor获取下游indata anchor，若下游算子的incontrol anchor与wait尚未连接，
        // 则与wait算子的outcontrol连接
        for (auto peerInDataAnchor : outDataAnchor->GetPeerInDataAnchors()) {
            if (!peerInDataAnchor) {
                continue;
            }
            auto ownerInControlAnchor = peerInDataAnchor->GetOwnerNode()->GetInControlAnchor();
            if (anchorSet.count((uintptr_t)ownerInControlAnchor.get()) == 0) {
                ge::GraphUtils::AddEdge(addedNodePtr->GetOutControlAnchor(), ownerInControlAnchor);
                anchorSet.insert((uintptr_t)ownerInControlAnchor.get());
            }
        }
        HCCL_DEBUG("outDataAnchor linkto PeerInDataAnchor num is [%d]", anchorSet.size());
    }
    return HCCL_SUCCESS;
}

HcclResult HvdGraphOptimizer::AddDependence(std::vector<ge::NodePtr> &preDependNodes)
{
    ge::graphStatus gRet;
    HCCL_DEBUG("preDependNodes size is [%d]", preDependNodes.size());
    if (preDependNodes.size() > 0) {
        HCCL_DEBUG("add dependence between HorovodWait");
        for (u32 i = 0; i < (preDependNodes.size() - 1); i++) {
            CHK_SMART_PTR_NULL(preDependNodes[i]->GetOutControlAnchor());
            CHK_SMART_PTR_NULL(preDependNodes[i + 1]->GetInControlAnchor());
            gRet = ge::GraphUtils::AddEdge(preDependNodes[i]->GetOutControlAnchor(),
                preDependNodes[i + 1]->GetInControlAnchor());
            bool bErr = (gRet != ge::GRAPH_SUCCESS);
            CHK_PRT_RET(bErr, HCCL_ERROR("[Add][Dependence]add preDependNodes dependence failed."), HCCL_E_INTERNAL);
        }
    }
    return HCCL_SUCCESS;
}

ge::Status HvdGraphOptimizer::OptimizeFusedGraph(ge::ComputeGraph& graph)
{
    for (auto nodePtr : graph.GetDirectNode()) {
        if (!nodePtr) {
            HCCL_WARNING("null node exists.");
            continue;
        }
        auto opDescPtr = nodePtr->GetOpDesc();
        if (!opDescPtr) {
            HCCL_WARNING("desc of node[%s] is null.", nodePtr->GetName().c_str());
            continue;
        }

        if (CheckSupportedOP(opDescPtr->GetType()) != HCCL_SUCCESS) {
            continue;
        }
        ge::Status ret = CalcOpRunningParam(*nodePtr);
        CHK_PRT_RET(ret != ge::SUCCESS, HCCL_ERROR("errNo[0x%016llx] Calc Op Running Params failed.",
            HCOM_ERROR_CODE(HCCL_E_INTERNAL)), ret);
    }
    return ge::SUCCESS;
}

ge::Status HvdGraphOptimizer::OptimizeWholeGraph(ge::ComputeGraph& graph)
{
    return ge::SUCCESS;
}

ge::Status HvdGraphOptimizer::GetAttributes(ge::GraphOptimizerAttribute& attrs) const
{
    attrs.engineName = HCCL_OPS_ENGIN;
    attrs.scope = ge::UNIT;
    HCCL_DEBUG("hccl graph optimizer get attr success. engine[%s] scope[%d]", attrs.engineName.c_str(), attrs.scope);
    return ge::SUCCESS;
}

ge::Status HvdGraphOptimizer::CalcOpRunningParam(ge::Node& node)
{
    HCCL_INFO("calculate hccl runing parameters start.");
    CHK_PRT_RET(!node.GetOpDesc(), HCCL_ERROR("[CalcOp][RunningParam]errNo[0x%016llx] GetOpDesc failed. null ptr.", \
        HCOM_ERROR_CODE(HCCL_E_PTR)), ge::INTERNAL_ERROR);

    std::string sCollectiveType = node.GetOpDesc()->GetType();
    HcclResult ret = CheckSupportedOP(sCollectiveType);
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[CalcOp][RunningParam]errNo[0x%016llx] op type[%s] is not supported.",
        HCOM_ERROR_CODE(ret), sCollectiveType.c_str()), ge::INTERNAL_ERROR);

    // No need to apply sub stream in OP_BASE mode
    if (ge::AttrUtils::SetInt(node.GetOpDesc(), "used_stream_num", 0) == false) {
        HCCL_ERROR("[CalcOp][RunningParam]errNo[0x%016llx] op[%s]: to OpDesc failed.", HCOM_ERROR_CODE(HCCL_E_PARA), \
            sCollectiveType.c_str());
        return ge::INTERNAL_ERROR;
    }
    // No need to set workspace mem in OP_BASE mode
    ret = SetOpOutputMemSize(node, sCollectiveType); // set output size
    CHK_PRT_RET(ret != HCCL_SUCCESS, HCCL_ERROR("[CalcOp][RunningParam]errNo[0x%016llx] set op[%s] output size failed.",
        HCOM_ERROR_CODE(ret), sCollectiveType.c_str()), ge::INTERNAL_ERROR);

    ret = SetOpMemAttr(node, sCollectiveType);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("[CalcOp][RunningParam]errNo[0x%016llx] set op[%s] continuous mem attr failed.",
            HCOM_ERROR_CODE(ret), sCollectiveType.c_str()), ge::INTERNAL_ERROR);

    ret = SetOpAtomicInputIndex(node, sCollectiveType);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("[CalcOp][RunningParam]errNo[0x%016llx] set op[%s] atomic input index failed.",
            HCOM_ERROR_CODE(ret),  sCollectiveType.c_str()), ge::INTERNAL_ERROR);

    HCCL_INFO("calculate hccl runing parameters completed.");
    return ge::SUCCESS;
}

HcclResult HvdGraphOptimizer::SetOpOutputMemSize(ge::Node& node, const std::string &sCollectiveType)
{
    ge::OpDescPtr op = node.GetOpDesc();
    for (u32 i = 0; i < op->GetOutputsSize(); i++) {
        int64_t memSize = 0;
        ge::GeTensorDesc outputTensor = op->GetOutputDesc(i);
        ge::GeShape outputShape = outputTensor.GetShape();
        ge::Format format = outputTensor.GetFormat();
        ge::DataType dataType = outputTensor.GetDataType();
        // 获取内存大小
        bool bErr = (ge::GRAPH_SUCCESS != ge::TensorUtils::CalcTensorMemSize(outputShape, format, dataType, memSize));
        CHK_PRT_RET(bErr, HCCL_ERROR("[Set][OutputMemSize]In get output mem size, error outputSize because no know"
            "shape, Format[%d], dataType[%d], outputSize[%lld], index[%u]", format,
            dataType, memSize, i), HCCL_E_PARA);

        if (memSize == -1) { // memsize 为-1 时，表示输入的shape不正确
            HCCL_ERROR("[Set][OutputMemSize]In get output mem size, error outputSize because unknow shape,"
                "Format[%d], dataType[%d], outputSize[%lld], index[%u]", format, dataType,
                memSize, i);
            return HCCL_E_PARA;
        }

        // 根据 规则重新计算内存大小
        HcclResult ret = CalcHCCLOutputMemSize(sCollectiveType, memSize);
        CHK_PRT_RET(ret != HCCL_SUCCESS,
            HCCL_ERROR("[Set][OutputMemSize]In calc output mem size, cacl, memsize error, ret[%d],"
            "sCollectiveType[%s], memSize[%lld]", ret, sCollectiveType.c_str(), memSize), HCCL_E_PARA);

        // 将内存大小重新传给上层
        ge::TensorUtils::SetSize(outputTensor, memSize);

        // 更新output Tensor
        if (op->UpdateOutputDesc(i, outputTensor)) {
            HCCL_ERROR("[Set][OutputMemSize]In get output mem size, update output desc error, Format[%d],"
                "dataType[%d], outputSize[%lld], index[%u]", format, dataType, memSize, i);
            return HCCL_E_PARA;
        }
        HCCL_INFO("In set output MemSize, sCollectiveType[%s], opMemSize[%lld]", sCollectiveType.c_str(), memSize);
    }
    return HCCL_SUCCESS;
}

HcclResult HvdGraphOptimizer::CalcHCCLOutputMemSize(const std::string &sCollectiveType, int64_t& memSize)
{
    const u32 MEMORY_ALIGN_RATIO = 2; // GE要求内存需要32KB对齐后，再外加32KB. out = (in + 2 * 32 - 1) / 32 * 32
    const u32 MEMORY_ALIGN_SIZE = 32; // GE要求内存需要32KB对齐后，再外加32KB. out = (in + 2 * 32 - 1) / 32 * 32
    // GE要求内存需要32KB对齐后，再外加32KB
    memSize = ((memSize + MEMORY_ALIGN_RATIO * MEMORY_ALIGN_SIZE - 1) / MEMORY_ALIGN_SIZE) * MEMORY_ALIGN_SIZE;
    return HCCL_SUCCESS;
}

HcclResult HvdGraphOptimizer::SetOpMemAttr(ge::Node& node, const std::string &sCollectiveType)
{
    bool bRet =  false;
    // broadcast op use the same buffer for input and output
    bool bRefAttr = (sCollectiveType == HVD_KERNEL_OP_TYPE_BROADCAST);
    bRet = ge::AttrUtils::SetBool(node.GetOpDesc(), ge::ATTR_NAME_REFERENCE, bRefAttr);
    CHK_PRT_RET(!bRet, HCCL_ERROR("[SetOp][MemAttr]errNo[0x%016llx] op[%s]: set  reference attr[%d] to OpDesc"
        "failed.", HCOM_ERROR_CODE(HCCL_E_PARA), sCollectiveType.c_str(), bRefAttr), HCCL_E_PARA);
    if (bRefAttr) {
        bRet = node.GetOpDesc()->UpdateOutputName(node.GetOpDesc()->GetAllInputName());
        CHK_PRT_RET(!bRet, HCCL_ERROR("[SetOp][MemAttr]errNo[0x%016llx] op[%s]: update output name failed.", \
            HCOM_ERROR_CODE(HCCL_E_PARA), sCollectiveType.c_str()), HCCL_E_PARA);
    }
    HCCL_INFO("hvd ops memory attr: op[%s], ref[%d]", sCollectiveType.c_str(), bRefAttr);

    return HCCL_SUCCESS;
}

HcclResult HvdGraphOptimizer::SetOpAtomicInputIndex(ge::Node& node, const std::string &sCollectiveType)
{
    // allreduce set atomic Input Index attribute
    if (sCollectiveType == HVD_KERNEL_OP_TYPE_ALLREDUCE) {
        vector<int64_t> atomicInputIndex(1, -1); // set value -1 as flag
        if (!ge::AttrUtils::SetListInt(node.GetOpDesc(), "atomic_input_index", atomicInputIndex)) {
            HCCL_ERROR("[SetOp][AtomicInputIndex]errNo[0x%016llx]: set op[%s] atomic index failed.",
                HCOM_ERROR_CODE(HCCL_E_PARA), sCollectiveType.c_str());
            return HCCL_E_PARA;
        }
    }
    return HCCL_SUCCESS;
}
}
